from typing import Callable, List, Tuple, Dict, Optional
import math
import numpy as np
from numpy.linalg import eigvalsh

from qiskit import QuantumCircuit, QuantumRegister
from qiskit.circuit.library import QFT
from qiskit.circuit import Qubit

import pytest

from qiskit.quantum_info import Statevector, Operator
rng = np.random.default_rng(2025)
from algorithm.qiskit.Hamiltonian_simulator import build_pauli_evolution_oracle_from_matrix, pauli_matrix

# evolution_oracle(qc, ctrl, wrk, tau): append controlled-U(tau) with control qubit `ctrl`
EvolutionOracle = Callable[[QuantumCircuit, Qubit, QuantumRegister, float], None]

# -----------------------------
# HHL PE helpers (measurement-free)
# -----------------------------

def hhl_pe_encode(
    qc: QuantumCircuit,
    phase_reg: QuantumRegister,
    work_reg: QuantumRegister,
    evolution_oracle: EvolutionOracle,
    t: float
):
    """H^m -> controlled-U(2^k t) -> IQFT, NO measurement."""
    for q in phase_reg:
        qc.h(q)
    qc.barrier()
    for k, ctrl in enumerate(phase_reg):
        power = 2 ** k
        evolution_oracle(qc, ctrl, work_reg, + power * t)
        qc.barrier()
    qc.append(QFT(len(phase_reg), inverse=True, do_swaps=False), phase_reg)

def hhl_pe_uncompute(
    qc: QuantumCircuit,
    phase_reg: QuantumRegister,
    work_reg: QuantumRegister,
    evolution_oracle: EvolutionOracle,
    t: float
):
    """QFT -> [controlled-U(2^k t)]^† -> H^m, to erase the phase register."""
    qc.append(QFT(len(phase_reg), inverse=False, do_swaps=False), phase_reg)
    for k, ctrl in reversed(list(enumerate(phase_reg))):
        power = 2 ** k
        evolution_oracle(qc, ctrl, work_reg, - power * t)
    for q in phase_reg:
        qc.h(q)

# -----------------------------
# t-chooser util
# -----------------------------

def pick_t_from_lambda_max(lambda_max: float, target_phase: float = 0.25, safety: float = 0.95) -> float:
    if lambda_max <= 0:
        raise ValueError("lambda_max must be > 0 (SPD assumed in HHL).")
    return safety * (2.0 * math.pi * target_phase) / lambda_max

def rough_lambda_max_via_eigs(H: np.ndarray) -> float:
    """For small matrices, just use eigs (dev utility). For large, replace with Gershgorin/power method."""
    w = eigvalsh(H)
    return float(np.max(np.abs(w)))

# test_hhl_pe_helpers.py
# -*- coding: utf-8 -*-

# -----------------------------
# 工具函数
# -----------------------------

def unitary_global_distance(U, V):
    """全局相位不敏感的 Frobenius 距离"""
    tr = np.trace(V.conj().T @ U)
    phi = np.angle(tr)
    return np.linalg.norm(U - np.exp(1j * phi) * V, ord='fro') / U.shape[0]


# -----------------------------
# 1) encode + uncompute 应为恒等（在 Suzuki2 下严格互逆）
# -----------------------------

@pytest.mark.parametrize("n", [1, 2])
@pytest.mark.parametrize("num_slices", [1, 2, 4])
def test_encode_uncompute_is_identity(n, num_slices):
    # 随机（但稳定）厄米 H：用 Pauli 和构造，含非对易项
    H = (0.7 * pauli_matrix("ZI"[:n]) +
         0.5 * (pauli_matrix(("X" * n)) if n == 1 else pauli_matrix("XY"[:n])) +
         0.3 * (pauli_matrix(("Y" * n)) if n == 1 else pauli_matrix("ZX"[:n])))
    tau = 0.37  # 任意

    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=num_slices)

    m = 3  # 相位寄存器位数
    phase = QuantumRegister(m, "ph")
    work  = QuantumRegister(n, "w")
    qc = QuantumCircuit(phase, work)

    # 做一次 encode + uncompute
    hhl_pe_encode(qc, phase, work, oracle, tau)
    qc.barrier()
    hhl_pe_uncompute(qc, phase, work, oracle, tau)
    U = Operator(qc).data
    I = np.eye(2**(m+n), dtype=complex)
    err = unitary_global_distance(U, I)
    assert err < 1e-12, f"encode+uncompute 未还原恒等（Suzuki2 下应严格互逆），误差={err:.3e}"


# -----------------------------
# 2) 相位读数：H = ω Z（1 比特工作寄存器）
#    选择 t 使 φ = ω t / (2π) = 1/4，3 位相位寄存器应得到 0.010（|0>）与 0.110（|1>）
# -----------------------------

def test_phase_readout_for_Z_eigenstates():
    n = 1
    omega = 1.0
    H = omega * pauli_matrix("Z")
    m = 3  # 3 位相位，可表示 k/8
    phi = 1/4  # 精确可表示（= 2/8）
    t = 2 * math.pi * phi / omega  # 保证 φ = ω t / (2π) = 1/4

    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=1)

    # |work> = |0>（Z 的 +1 本征态），phase 初始 |000>
    ph = QuantumRegister(m, "ph")
    w  = QuantumRegister(n, "w")
    qc0 = QuantumCircuit(ph, w)

    hhl_pe_encode(qc0, ph, w, oracle, t)
    sv0 = Statevector.from_instruction(qc0)

    # 只对相位寄存器做边缘概率（返回 MSB-first 的 bitstring）
    phase_indices = [qc0.find_bit(q).index for q in ph]
    probs0 = sv0.probabilities_dict(qargs=phase_indices)
    bits0 = max(probs0, key=probs0.get)   # 最可能的相位位串
    assert bits0 == '010'  # φ = 1/4

    # |work> = |1>（Z 的 -1 本征态），相位 3/4 -> |110>
    qc1 = QuantumCircuit(ph, w)
    qc1.x(w[0])
    hhl_pe_encode(qc1, ph, w, oracle, t)
    sv1 = Statevector.from_instruction(qc1)
    phase_indices = [qc1.find_bit(q).index for q in ph]
    probs1 = sv1.probabilities_dict(qargs=phase_indices)
    bits1 = max(probs1, key=probs1.get)
    assert bits1 == '110'  # φ = 3/4


# -----------------------------
# 3) 含 c_I I 的情形：encode+uncompute 仍应为恒等
# -----------------------------

def test_identity_term_does_not_break_uncompute():
    n = 2
    # H = c_I I + 其他 Pauli
    c_I = 1.234
    H = (c_I * np.eye(2**n, dtype=complex)
         + 0.6 * pauli_matrix("XZ")
         - 0.4 * pauli_matrix("ZX"))
    t = 0.51
    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=2)

    m = 3
    ph = QuantumRegister(m, "ph")
    w  = QuantumRegister(n, "w")
    qc = QuantumCircuit(ph, w)

    hhl_pe_encode(qc, ph, w, oracle, t)
    hhl_pe_uncompute(qc, ph, w, oracle, t)

    U = Operator(qc).data
    I = np.eye(2**(m+n), dtype=complex)
    err = unitary_global_distance(U, I)
    assert err < 1e-12, f"含 c_I 时 encode+uncompute 未完全抵消，误差={err:.3e}"


# -----------------------------
# 4) t 选择函数 & λ_max 估计 sanity check
# -----------------------------

def test_pick_t_and_lambda_max():
    # SPD 情况（HHL 假设）
    H = 0.2 * pauli_matrix("Z") + 0.3 * pauli_matrix("I")  # λ_max = 0.5
    lam = rough_lambda_max_via_eigs(H)
    assert 0.4 < lam < 0.6

    t = pick_t_from_lambda_max(lam, target_phase=0.25, safety=0.9)
    # 期望数量级：t ≈ 0.9 * (2π * 0.25) / 0.5 = 0.9 * π ≈ 2.827
    assert 2.5 < t < 3.2, f"pick_t 返回值数量级异常：t={t}"

    with pytest.raises(ValueError):
        pick_t_from_lambda_max(0.0)